import os
from torch.utils.data import Dataset
import pandas as pd
from PIL import Image
from torchvision import transforms
from data import custom_transforms as tr
from torchvision import utils as vutils
import matplotlib.pyplot as plt

class DataProcessor():
    def __init__(self,base_size=64):
        self.base_size=base_size

    def process_data(self,data_path):
        # get all the img paths in data
        data_path = os.path.abspath(data_path)
        imgs_list = []
        root_dirs = os.listdir(data_path)
        for dir in root_dirs:
            road_path = os.path.join(data_path, dir)
            if dir.startswith("Road") and os.path.isdir(road_path):
                color_img_path = os.path.join(road_path, "ColorImage_" + dir.lower(), "ColorImage")
                record_dirs = os.listdir(color_img_path)
                for record_dir in record_dirs:
                    record_path = os.path.join(color_img_path, record_dir)
                    camera_dirs = os.listdir(record_path)
                    for camera_dir in camera_dirs:
                        camera_path = os.path.join(record_path, camera_dir)
                        imgs = os.listdir(camera_path)
                        for img in imgs:
                            img_path = os.path.join(camera_path, img)
                            if os.path.isfile(img_path):
                                imgs_list.append(img_path)
        for img_path in imgs_list:
            input_img = Image.open(img_path)
            dirs= img_path.split(os.sep)
            label_path= os.path.join(data_path,"Gray_Label", 'Label_'+dirs[-6].lower(),"Label",dirs[-3],dirs[-2],dirs[-1].replace(".jpg","_bin.png"))
            label = Image.open(label_path)
            sample = {'image': input_img, 'label': label}
            imgs = self.transform_tr(sample)
            # vutils.save_image(imgs["image"], os.path.join(data_path, "processed_data","input",img_path.split(os.sep)[-1]))
            # print()
            # output_img = transforms.ToPILImage()(imgs["image"])
            # output_label = transforms.ToPILImage()(imgs["label"])
            # plt.imshow(output_img)
            # plt.show()
            imgs["image"].save(os.path.join(data_path, "processed_data","input",img_path.split(os.sep)[-1]))
            imgs["label"].save(os.path.join(data_path, "processed_data","label",label_path.split(os.sep)[-1]))

    def transform_tr(self,sample):
        composed_transforms = transforms.Compose([
            tr.FixedResize(256,512),
            # tr.RandomHorizontalFlip(),
            # tr.RandomScaleCrop(base_size=self.base_size, crop_size=self.base_size),
            tr.RandomGaussianBlur(),
            # tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            # tr.ToTensor()
        ])
        return composed_transforms(sample)
if __name__ == '__main__':
    dp = DataProcessor()
    dp.process_data("../../RoadLineDataset")
